'''
DeepLabV3Plus-Pytorch ONNX推理 - ResNet101 Cityscapes
----------------------------------------------
@作者: gx
@邮箱: gaoxukkk@qq.com
@创建日期：2024年11月14日
@修改日期：2024年11月14日 - ONNX推理脚本
'''

import cv2
import numpy as np
import onnxruntime as ort
from PIL import Image
import os
import time

# Cityscapes颜色映射表
def cityscapes_cmap():
    """
    Cityscapes数据集的颜色映射表
    返回19个类别的RGB颜色值
    """
    # Cityscapes的19个类别颜色映射 (不包括ignore的类别)
    colors = [
        [128, 64, 128],   # road
        [244, 35, 232],   # sidewalk
        [70, 70, 70],     # building
        [102, 102, 156],  # wall
        [190, 153, 153],  # fence
        [153, 153, 153],  # pole
        [250, 170, 30],   # traffic light
        [220, 220, 0],    # traffic sign
        [107, 142, 35],   # vegetation
        [152, 251, 152],  # terrain
        [70, 130, 180],   # sky
        [220, 20, 60],    # person
        [255, 0, 0],      # rider
        [0, 0, 142],      # car
        [0, 0, 70],       # truck
        [0, 60, 100],     # bus
        [0, 80, 100],     # train
        [0, 0, 230],      # motorcycle
        [119, 11, 32],    # bicycle
    ]
    return np.array(colors, dtype=np.uint8)

# Cityscapes类别名称
def get_cityscapes_class_names():
    """
    返回Cityscapes的类别名称
    """
    class_names = [
        'road', 'sidewalk', 'building', 'wall', 'fence',
        'pole', 'traffic light', 'traffic sign', 'vegetation', 'terrain',
        'sky', 'person', 'rider', 'car', 'truck',
        'bus', 'train', 'motorcycle', 'bicycle'
    ]
    return class_names

def preprocess_image(image_path, target_size=(1024, 2048)):
    """
    预处理图像
    Args:
        image_path: 输入图像路径
        target_size: 目标尺寸 (height, width)
    Returns:
        预处理后的图像tensor
    """
    # 读取图像
    image = cv2.imread(image_path)
    if image is None:
        raise ValueError(f"Cannot read image from {image_path}")
    
    # 调整尺寸
    image = cv2.resize(image, (target_size[1], target_size[0]))
    
    # BGR to RGB
    image = image[..., ::-1]
    
    # 归一化到[0, 1]
    image = (image / 255.0).astype(np.float32)
    
    # 标准化 (ImageNet mean/std)
    mean = [0.485, 0.456, 0.406]
    std = [0.229, 0.224, 0.225]
    for c in range(image.shape[2]):
        image[..., c] = (image[..., c] - mean[c]) / std[c]
    
    # 转换为CHW格式
    image = image.transpose(2, 0, 1)
    
    # 添加batch维度
    image = image[None, ...]
    
    return image

def postprocess_output(output, original_size=None):
    """
    后处理模型输出
    Args:
        output: 模型输出 (1, 19, H, W)
        original_size: 原始图像尺寸 (H, W)
    Returns:
        分割结果图像
    """
    # 获取预测类别
    pred = np.argmax(output[0], axis=0)  # (H, W)
    
    # 颜色映射
    cmap = cityscapes_cmap()
    colorized_pred = cmap[pred].astype(np.uint8)
    
    # 如果需要，调整回原始尺寸
    if original_size is not None:
        colorized_pred = cv2.resize(colorized_pred, (original_size[1], original_size[0]))
    
    return colorized_pred, pred

def create_overlay(original_image, segmentation_result, alpha=0.7):
    """
    创建叠加图像
    Args:
        original_image: 原始图像 (BGR格式)
        segmentation_result: 分割结果 (RGB格式)
        alpha: 透明度
    Returns:
        叠加后的图像
    """
    # 转换分割结果为BGR格式
    seg_bgr = cv2.cvtColor(segmentation_result, cv2.COLOR_RGB2BGR)
    
    # 叠加
    overlay = cv2.addWeighted(original_image, 1-alpha, seg_bgr, alpha, 0)
    
    return overlay

def main():
    # 配置参数
    onnx_model_path = "checkpoints/deeplabv3plus_resnet101_cityscapes.onnx"
    input_image_path = "/data2/xd/DeepLabV3Plus-Pytorch/samples/undistorted_image.jpg"
    output_dir = "results"
    
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    
    # 检查ONNX模型文件是否存在
    if not os.path.exists(onnx_model_path):
        print(f"Error: ONNX model not found at {onnx_model_path}")
        print("Please run export_onnx_resnet101_cityscapes.py first to generate the ONNX model.")
        return
    
    # 检查输入图像是否存在
    if not os.path.exists(input_image_path):
        print(f"Error: Input image not found at {input_image_path}")
        return
    
    print(f"Loading ONNX model from {onnx_model_path}")
    
    # 创建ONNX推理会话
    providers = ['CPUExecutionProvider']  # 使用CPU推理，如果有GPU可以改为'CUDAExecutionProvider'
    session = ort.InferenceSession(onnx_model_path, providers=providers)
    
    # 获取输入输出信息
    input_name = session.get_inputs()[0].name
    output_name = session.get_outputs()[0].name
    input_shape = session.get_inputs()[0].shape
    output_shape = session.get_outputs()[0].shape
    
    print(f"Model input shape: {input_shape}")
    print(f"Model output shape: {output_shape}")
    
    # 读取原始图像用于叠加
    original_image = cv2.imread(input_image_path)
    original_size = original_image.shape[:2]  # (H, W)
    
    print(f"Processing image: {input_image_path}")
    print(f"Original image size: {original_size}")
    
    # 预处理图像
    start_time = time.time()
    input_tensor = preprocess_image(input_image_path)
    preprocess_time = time.time() - start_time
    
    print(f"Preprocessing time: {preprocess_time:.3f}s")
    
    # 推理
    start_time = time.time()
    outputs = session.run([output_name], {input_name: input_tensor})
    inference_time = time.time() - start_time
    
    print(f"Inference time: {inference_time:.3f}s")
    print(f"FPS: {1.0/inference_time:.2f}")
    
    # 后处理
    start_time = time.time()
    segmentation_result, pred_mask = postprocess_output(outputs[0], original_size)
    postprocess_time = time.time() - start_time
    
    print(f"Postprocessing time: {postprocess_time:.3f}s")
    
    # 创建叠加图像
    overlay_result = create_overlay(original_image, segmentation_result)
    
    # 保存结果
    base_name = os.path.splitext(os.path.basename(input_image_path))[0]
    
    # 保存分割结果
    seg_output_path = os.path.join(output_dir, f"{base_name}_segmentation.png")
    cv2.imwrite(seg_output_path, cv2.cvtColor(segmentation_result, cv2.COLOR_RGB2BGR))
    print(f"Segmentation result saved to: {seg_output_path}")
    
    # 保存叠加结果
    overlay_output_path = os.path.join(output_dir, f"{base_name}_overlay.png")
    cv2.imwrite(overlay_output_path, overlay_result)
    print(f"Overlay result saved to: {overlay_output_path}")
    
    # 保存预测掩码（灰度图）
    mask_output_path = os.path.join(output_dir, f"{base_name}_mask.png")
    mask_resized = cv2.resize(pred_mask.astype(np.uint8), (original_size[1], original_size[0]))
    cv2.imwrite(mask_output_path, mask_resized)
    print(f"Prediction mask saved to: {mask_output_path}")
    
    # 统计类别分布
    class_names = get_cityscapes_class_names()
    unique_classes, counts = np.unique(pred_mask, return_counts=True)
    
    print("\nClass distribution:")
    for class_id, count in zip(unique_classes, counts):
        if class_id < len(class_names):
            percentage = (count / pred_mask.size) * 100
            print(f"  {class_names[class_id]}: {count} pixels ({percentage:.2f}%)")
    
    print(f"\nTotal processing time: {preprocess_time + inference_time + postprocess_time:.3f}s")
    print("Inference completed successfully!")

if __name__ == "__main__":
    main() 